import torch
from transformers import BertTokenizer, BertModel
from models.sentiment_model import SentimentClassifier
import pandas as pd


def load_model(model_path):
    # 设置weights_only=False以完整加载模型文件
    checkpoint = torch.load(model_path, weights_only=False)

    bert_model = BertModel.from_pretrained('bert-base-chinese')
    model = SentimentClassifier(
        n_classes=checkpoint['config']['n_classes'],
        pretrained_model=bert_model
    )

    model.load_state_dict(checkpoint['model_state_dict'])
    tokenizer = checkpoint['tokenizer']
    max_len = checkpoint['config']['max_len']

    return model, tokenizer, max_len

def predict_sentiment(text, model, tokenizer, max_len):
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model = model.to(device)
    model.eval()
    
    encoding = tokenizer.encode_plus(
        text,
        add_special_tokens=True,
        max_length=max_len,
        padding='max_length',
        truncation=True,
        return_tensors='pt'
    )
    
    input_ids = encoding['input_ids'].to(device)
    attention_mask = encoding['attention_mask'].to(device)
    
    with torch.no_grad():
        outputs = model(input_ids=input_ids, attention_mask=attention_mask)
        prediction = torch.argmax(outputs, dim=1)
    
    return prediction.item()

def test_single_text(model_path):
    model, tokenizer, max_len = load_model(model_path)
    while True:
        print("\n请输入要测试的文本（输入'q'退出）:")
        text = input().strip()
        if text.lower() == 'q':
            break
            
        pred_label = predict_sentiment(text, model, tokenizer, max_len)
        sentiment = "积极" if pred_label == 1 else "消极"
        print(f"文本: {text}")
        print(f"情感预测: {sentiment}")

def test_from_file(test_file, model_path):
    model, tokenizer, max_len = load_model(model_path)
    df = pd.read_csv(test_file, encoding='utf-8')
    
    correct = 0
    total = len(df)
    results = []
    
    print(f"开始测试，共{total}条数据...")
    
    for idx, row in enumerate(df.iterrows(), 1):
        text = row[1]['review']
        true_label = row[1]['label']
        pred_label = predict_sentiment(text, model, tokenizer, max_len)
        
        is_correct = pred_label == true_label
        if is_correct:
            correct += 1
            
        results.append({
            'text': text,
            'true_label': true_label,
            'predicted_label': pred_label,
            'is_correct': is_correct
        })
        
        # 每100条数据显示一次进度
        if idx % 100 == 0:
            current_acc = correct / idx
            print(f"已处理: {idx}/{total} 条数据，当前准确率: {current_acc:.4f}")
    
    final_accuracy = correct / total
    print(f'\n测试完成！')
    print(f'总体准确率: {final_accuracy:.4f}')
    
    # 保存详细结果
    results_df = pd.DataFrame(results)
    results_df.to_csv('test_results.csv', index=False, encoding='utf-8')
    print(f'详细测试结果已保存到 test_results.csv')

if __name__ == '__main__':
    model_path = 'models/sentiment_model/best_model.pth'
    test_file = 'data/test_generated.csv'  # 使用生成的测试数据

    print(f"使用设备: {torch.device('cuda' if torch.cuda.is_available() else 'cpu')}")
    
    while True:
        print("\n请选择测试模式：")
        print("0: 单条测试")
        print("1: 批量测试")
        print("q: 退出程序")
        
        choice = input().strip().lower()
        if choice == 'q':
            break
            
        try:
            mode = int(choice)
            if mode == 0:
                test_single_text(model_path)
            elif mode == 1:
                test_from_file(test_file, model_path)
            else:
                print("无效的选择，请输入0或1")
        except ValueError:
            print("无效的输入，请输入0或1")